﻿// Left-leaning red-black trees a la Sedgewick 2008/Andersson 1993
// sestoft@itu.dk * 2011-03-28, 2011-03-31

// Paper at: http://www.cs.princeton.edu/~rs/talks/LLRB/LLRB.pdf
// Some source at: http://www.cs.princeton.edu/~rs/talks/LLRB/Java/

// Warning: It seems to be an assumption of Sedgewick's "remove" code
// that the item to be removed is actually in the tree -- else trivial
// null pointer exception on null childs.  Fixed in this code.

// For testing, run with 
//   java -ea LeftLeaning

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;
using C5;
using System.Diagnostics.Contracts;
using snapshottrees;

namespace snapshottrees
{
    public class LLTree : ITree
    {
        // A node carries the color of the link to it.  Every leaf is black.
        // Invariant: No red node has a red child.
        // Invariant: Every leaf path has the same number of black nodes.
        // Invariant: Only the left child can be red (left-leaning 2-3 tree).
        [ContractInvariantMethod]
        void ObjectInvariant()
        {
            Contract.Invariant(!isRed(this.root));
            Contract.Invariant(Math.Pow(2.0, check(root)) <= count(root) + 1);
            //Contract.Invariant(Contract.ForAll(getAllNodes(root), node => !isRed(node.rght)));
            //Contract.Invariant(Contract.ForAll(getAllNodes(root), node => isRed(node) ? !isRed(node.left) : isRed(node.left) || !isRed(node.left)));
        }


        private Node root;            // Is null for empty set
        public static readonly Boolean RED = true, BLACK = false;

        public class Node
        {
            public int item;
            public Node left;
            public Node rght;
            public Boolean color;

            public Node(int item)
            {
                Contract.Ensures(this.item == item);
                Contract.Ensures(this.color == RED);

                this.item = item;
                this.color = RED;
            }

            public void flipColors()
            {
                Contract.Ensures(this == null || this.color != Contract.OldValue<Boolean>(this.color));
                Contract.Ensures(this.left == null || this.left.color != Contract.OldValue<Boolean>(this.left.color));
                Contract.Ensures(this.rght == null || this.rght.color != Contract.OldValue<Boolean>(this.rght.color));
                
                color = !color;
                left.color = !left.color;
                rght.color = !rght.color;
            }

            public Node rotateLeft()
            {
                Contract.Ensures(isRed(this));
                Contract.Ensures(this.rght == Contract.OldValue<Node>(this.rght.left));
                Contract.Ensures(Contract.Result<Node>().left == this);
                Contract.Ensures(Contract.Result<Node>().rght == Contract.OldValue<Node>(this.rght.rght));
                Contract.Ensures(this.left == Contract.OldValue<Node>(this.left));
                Contract.Ensures(Contract.Result<Node>().color == Contract.OldValue<Boolean>(this.color));

                Node x = rght;
                rght = x.left;
                x.left = this;
                x.color = color;
                color = RED;
                return x;
            }

            public Node rotateRght()
            {
                Contract.Ensures(isRed(this));
                Contract.Ensures(this.left == Contract.OldValue<Node>(this.left.rght));
                Contract.Ensures(Contract.Result<Node>().rght == this);
                Contract.Ensures(Contract.Result<Node>().left == Contract.OldValue<Node>(this.left.left));
                Contract.Ensures(this.rght == Contract.OldValue<Node>(this.rght));
                Contract.Ensures(Contract.Result<Node>().color == Contract.OldValue<Boolean>(this.color));
                
                Node x = left;
                left = x.rght;
                x.rght = this;
                x.color = color;
                color = RED;
                return x;
            }
            [Pure]
            public String toString()
            {
                Contract.Ensures(Contract.Result<String>() != null);

                return (color==RED ? "R" : "B")
                    + "[" + (left==null ? "_" : left.toString()) + "," 
                    + item + "," + (rght==null ? "_" : rght.toString()) + "]";
            }
        }

        public LLTree()
        {
        }

        [Pure]
        public Boolean contains(int item)
        {
            Node node = root;
            Boolean found = false;
            while (!found && node != null) 
                if (item < node.item)
                    node = node.left;
                else if (node.item < item)
                    node = node.rght;
                else
                    found = true;
            return found;
        }

        public Boolean add(int item)
        {
            Contract.Ensures(!isRed(root));
            
            RefBool updated = new RefBool();
            root = addRecursive(root, item, updated);
            Contract.Assume(this.contains(item));
            
            root.color = BLACK;
            return updated.value;
        }

        private Node addRecursive(Node node, int item, RefBool updated)
        {
            Contract.Requires(updated != null);
            Contract.Ensures(Contract.Result<Node>() != null);
            Contract.Ensures(Contract.Result<Node>().left == null || Contract.Result<Node>().left.item < Contract.Result<Node>().item);
            Contract.Ensures(Contract.Result<Node>().rght == null || Contract.Result<Node>().rght.item > Contract.Result<Node>().item);
            
            Node res = node;
            if (node == null)
            {
                updated.value = true;
                res = new Node(item);
            }
            else
            {
                if (item < node.item)
                {
                    Node newLeft = addRecursive(node.left, item, updated);
                    node.left = newLeft;    
                }
                else if (node.item < item)
                {
                    Node newRght = addRecursive(node.rght, item, updated);
                    node.rght = newRght;
                } // else item == node.item so no update.
                res = fixUp(res);
            }
            return res;
        }

        private static Node fixUp(Node node)
        {
            Contract.Ensures(Contract.Result<Node>().left == null || Contract.Result<Node>().left.item < Contract.Result<Node>().item);
            Contract.Ensures(Contract.Result<Node>().rght == null || Contract.Result<Node>().rght.item > Contract.Result<Node>().item);
            
            // Rebalance as left-leaning red-black (2-3) tree:
            if (isRed(node.rght) && !isRed(node.left))
                node = node.rotateLeft();
            if (isRed(node.left) && isRed(node.left.left))
                node = node.rotateRght();
            if (isRed(node.left) && isRed(node.rght))
                node.flipColors();
            return node;
        }
        
        [ContractVerification (true)]
        [Pure]
        public static Boolean isRed(Node node)
        {
            Contract.Ensures(Contract.Result<Boolean>() && node != null && node.color == RED
                         || !Contract.Result<Boolean>() && node == null
                         || !Contract.Result<Boolean>() && node.color != RED);

            return node != null && node.color==RED;  // A leaf is black
        }

        public Boolean remove(int item)
        {
            Contract.Ensures(!isRed(root));
            
            RefBool updated = new RefBool();    
            root = removeRecursive(root, item, updated);
            Contract.Assume(!this.contains(item));

            if (root != null)
                root.color = BLACK;
            return updated.value;
        }

        private Node removeRecursive(Node node, int item, RefBool updated)
        {
            Contract.Requires(updated != null);
            
            Node res = node;
            if (res != null)
            {
                if (item < res.item)
                {
                    if (res.left != null && !isRed(res.left) && !isRed(res.left.left))
                    res = moveRedLeft(res);
                    res.left = removeRecursive(res.left, item, updated);
                }
                else
                {
                    if (isRed(res.left))
                        res = res.rotateRght();
                    if (res.rght != null && !isRed(res.rght) && !isRed(res.rght.left))
                        res = moveRedRght(res);
                    if (item == res.item)
                    {
                        updated.value = true;
                        if (res.rght == null) 
                            res = null;
                        else
                        {
                            res.item = min(res.rght).item;
                            res.rght = removeMin(res.rght);
                        }
                    }
                    else // item > res.item
                        res.rght = removeRecursive(res.rght, item, updated);
                }
                if (res != null)
                    res = fixUp(res);
            }
            return res;
        }

        public void removeMin()
        {
            Contract.Ensures(!LLTree.isRed(root));

            root = removeMin(root);
            root.color = BLACK;
        }

        private Node removeMin(Node node)
        {
            Contract.Ensures((Contract.Result<Node>() == null && (node == null || node.left == null))
                          || (Contract.Result<Node>() != null && node != null));
            
            Node res = node;
            if (res == null || res.left == null)
                return null;
            if (!isRed(res.left) && !isRed(res.left.left))
                res = moveRedLeft(res);
            res.left = removeMin(res.left);
            res = fixUp(res);
            return res;
        }

        private Node moveRedLeft(Node node)
        {
            Contract.Requires(node != null);
            Contract.Ensures(Contract.Result<Node>() != null);
           
            Node res = node;
            res.flipColors();
            if (isRed(res.rght.left))
            {
                res.rght = res.rght.rotateRght();
                res = res.rotateLeft();
                res.flipColors();
            }
            return res;
        }

        private Node moveRedRght(Node node)
        {
            Contract.Requires(node != null);
            Contract.Ensures(Contract.Result<Node>() != null);

            Node res = node;
            res.flipColors();
            if (isRed(res.left.left))
            {
                res = res.rotateRght();
                res.flipColors();
            }
            return res;
        }

        private Node min(Node node)
        {
            Contract.Requires(node != null);
            Contract.Ensures(Contract.Result<Node>() != null && Contract.Result<Node>().left == null);

            Node res = node;
            while (res.left != null)
                res = res.left;
            return res;  // res != null && res.left == null
        }

        // Utilities for checking tree consistency etc

        [Pure]
        private static int count(Node n)
        {
            Contract.Ensures(Contract.Result<int>() == 0 && n == null
                          || Contract.Result<int>() == count(n.left) + 1 + count(n.rght) && n != null);

            int res;
            if (n == null)
                res = 0;
            else 
                res = count(n.left) + 1 + count(n.rght);
            return res;
        }

        [Pure]
        private static int depth(Node n)
        {
            Contract.Ensures(Contract.Result<int>() == 0 && n == null
                          || Contract.Result<int>() == 1 + Math.Max(depth(n.left), depth(n.rght)) && n != null);

            if (n == null)
                return 0;
            else 
                return 1 + Math.Max(depth(n.left), depth(n.rght));
        }

        [Pure]
        public int check()
        {
            int blackDepth = check(root);
            Contract.Assert(Math.Pow(2.0, blackDepth) <= count(root) + 1, "tree unbalanced " + root.toString());
            return blackDepth;
        }

        // Return black-depth if tree well-formed, else throw exception
        [Pure]
        private static int check(Node node)
        {
            Contract.Ensures(Contract.Result<int>() == 0 && node == null
                          || Contract.Result<int>() >= 0 && node != null);

            if (node == null) 
                return 0;  
            else
            {
                if (isRed(node) && isRed(node.left))
                    throw new Exception("red node has red child: " + node.toString());
                if (isRed(node.rght))
                    throw new Exception("right child is red: " + node.toString());
                int blackDepth = check(node.left);
                if (blackDepth != check(node.rght))
                    throw new Exception("black depths differ: " + node.toString());	
                else
                    return blackDepth + (isRed(node) ? 0 : 1);
            }
        }

        [Pure]
        public void printTree()
        {
            int blackDepth = check();
            if (root != null)
                Console.WriteLine(root);
            Console.WriteLine("Blackdepth = %d; Depth = %d%n", blackDepth, depth(root));
        }

        //[Pure]
        //private static List<Node> getAllNodes(Node node)
        //{
        //    Contract.Requires(node != null);
        //    Contract.Ensures(Contract.Result<List<Node>>().Contains<Node>(node));
        //    Contract.Ensures(Contract.Result<List<Node>>() != null);

        //    List<Node> nodeList = new List<Node>();
        //    nodeList.Add(node);

        //    if (node.left != null)
        //        nodeList.AddRange(getAllNodes(node.left));
        //    if (node.rght != null)
        //        nodeList.AddRange(getAllNodes(node.rght));

        //    Contract.Assume(nodeList.Contains<Node>(node));
        //    return nodeList;
        //}
    }

    // For returning a boolean via a "ref" parameter
    public class RefBool
    {
        public Boolean value;
    }

    public interface ITree
    {
        Boolean contains(int item);
        Boolean add(int item);
        Boolean remove(int item);
        int check();
        void printTree();
    }

    public class Program
    {
        public static void Main(String[] args)
        {

        }
    }
}
